package org.skipper.dubbo.spring.starter.helper;

import lombok.extern.log4j.Log4j2;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.*;
import org.skipper.dubbo.spring.starter.model.RpcTraceContext;

import java.util.HashMap;
import java.util.Map;

import static org.apache.dubbo.common.constants.CommonConstants.CONSUMER;
import static org.apache.dubbo.common.constants.CommonConstants.PROVIDER;

/**
 * Created with IntelliJ IDEA.
 * Description:
 *
 * @author shoujian
 * Date: 2022-02-14
 * Time: 8:24 PM
 */
@Activate(group = {CONSUMER})
@Log4j2
public class DubboConsumerContextFilter implements Filter {
    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {

        Map<String, Object> attachments = RpcContext.getContext().getObjectAttachments();
        try {
            Map<String, Object> objectAttachments = RpcTraceContext.getTraceContext();

            if (objectAttachments != null) {

                for (String key : objectAttachments.keySet()) {
                    if (objectAttachments.get(key) == null) {
                        continue;
                    }
                    attachments.put(key, objectAttachments.get(key).toString());
                }
            }
            return invoker.invoke(invocation);
        }finally {
            RpcContext.getContext().setObjectAttachments(attachments);
        }
    }
}
